{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1be69dd2-b703-459e-980c-b91e34784402",
   "metadata": {},
   "source": [
    "<img alt=\"[GoMLX Mascot]\" src=\"gomlx_gopher2.jpg\" style=\"height:10em; float:left; margin:1em;\"/>\n",
    "\n",
    "# GoMLX Tutorial \n",
    "\n",
    "If you want just to quickly look at an working example, checkout [examples/cifar/demo/adult.ipynb](https://github.com/gomlx/gomlx/blob/main/examples/cifar/cifar.ipynb), for model trained on the UCI Adult Income dataset. More examples in [examples/](https://github.com/gomlx/gomlx/tree/main/examples) subdirectory. \n",
    "\n",
    "The tutorial won't detail the whole API, but should present all important concepts. Everything else is\n",
    "well documented in godoc (in the code), also available in [pkg.go.dev](https://pkg.go.dev/github.com/gomlx/gomlx#section-readme).\n",
    "\n",
    "The tutorial was written using a Jupyter notebook with [GoNB](https://github.com/janpfeifer/gonb), a kernel for Go that was co-developed with **GoMLX**. It has its own [short tutorial](https://github.com/janpfeifer/gonb/blob/main/examples/tutorial.ipynb) for those interested. But **GoMLX** doesn't require **GoNB**.\n",
    "\n",
    "If you are seeing this tutorial from [github](https://github.com/gomlx/gomlx/blob/main/examples/tutorial/tutorial.ipynb) snapshop, you won't be able to interact with it. To be able to play with it, try **installing GoMLX**, see its [README Installation section](https://github.com/gomlx/gomlx#installation). The easiest way is to start the pre-generated docker and use the Jupyter notebook there -- this tutorial can be opened from there in an interactive way.\n",
    "\n",
    "## Note: Output Not Displaying in JupyterLab ?\n",
    "\n",
    "If your notebook plots are not displaying correctly, it's likely because Jupyter is assuming the notebook is \"not trusted\". This can be easily fixed by opening the command palette (`Shift+Control+C` in PCs) and selecting `Trust Notebook`.\n",
    " There is an [issue opened in JupyterLab](https://github.com/jupyterlab/jupyterlab/issues/15467) attempting to address that.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a8d8449f-d405-4ec5-a69d-3aacfbcedbec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t- Replace rule for module \"github.com/gomlx/gomlx\" to local directory \"/home/janpf/Projects/gomlx\" already exists.\n"
     ]
    }
   ],
   "source": [
    "!*rm -f go.work && go work init && go work use . \"${HOME}/Projects/gomlx\"\n",
    "%goworkfix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53d4a4b5-7d5b-4293-b9d4-88bb00c26abf",
   "metadata": {},
   "source": [
    "## Computation Graphs (or Symbolic Computation)\n",
    "\n",
    "> [Package `graph` reference documentation](https://pkg.go.dev/github.com/gomlx/gomlx/pkg/core/graph)\n",
    "\n",
    "To do machine learning based on neural networks and gradient descent one of the most important\n",
    "requirements is the ability to do mathematical computations (mostly matrix multiplications)\n",
    "fast.\n",
    "\n",
    "GoMLX is built on the concept of building \"computation graphs\", just-in-time compiling them \n",
    "and only then executing them to get the desired results. That means one has to write code\n",
    "that generates other type of code (computation graph) so to say. We do this because then we are\n",
    "able to execute it really fast using [XLA](https://openxla.org/) using various accelerators (GPU, TPUs, etc).\n",
    "\n",
    "For example, let's create a computation graph to sum two values:\n",
    "\n",
    "> **Note**\n",
    "> - If executing this on a notebook, notice the very first cell execution takes a few seconds for Go to fetch the required dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e6c58651-25bb-47ab-ab60-6e8a7beeebe6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import . \"github.com/gomlx/gomlx/pkg/core/graph\"\n",
    "\n",
    "func Sum(a, b *Node) *Node {\n",
    "    return Add(a, b)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2631ec80-f864-4743-8124-757f430ed221",
   "metadata": {},
   "source": [
    "> **Note**\n",
    "> \n",
    "> - The `import . \"github.com/gomlx/gomlx/pkg/core/graph\"` import all definitions in computation to the current scope. Usually,\n",
    ">   it's easier to work like this on go files that are going to implement graph building functions.\n",
    "> - The type `*Node` represents a node in the computation graph. All graph operations either take a\n",
    ">   `*Graph` object to start with, or a `*Node`, and create new nodes with the corresponding operations.\n",
    ">   So our example will create an `Add` node that will take the nodes pointed by `a` and `b`, build a node \n",
    ">   that represent their summation and then return this `*Node`. \n",
    "> - Every node contains a reference to the `*Graph` it's part of (see [`Node.Graph()`](https://pkg.go.dev/github.com/gomlx/gomlx/pkg/core/graph#Node.Graph)).\n",
    "> - There is a rich set of operations available in GoMLX, see [`Node` documentation](https://pkg.go.dev/github.com/gomlx/gomlx/pkg/core/graph#Node).\n",
    "\n",
    "Ok, but this is \"symbolic\" only, it won't tell us what is 1+1 yet. We need to compile and then execute this graph with\n",
    "some input values."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41235bcc-8acb-478a-86ae-bd67f9ab0766",
   "metadata": {},
   "source": [
    "## Backends and executing Graphs with `graph.Exec`\n",
    "\n",
    "**GoMLX** can work with different backends to execute its operations. Currently 2 backends are supported:\n",
    "\n",
    "1. `\"xla\"`: It uses [OpenXLA/XLA](https://openxla.org/), a \"state of the art\" execution engine for ML operations, that supports CPU (with SIMD), various accelerators (GPUs, TPUs), JIT (Just In Time) compilation, providing the best performance. For now though, it only works in Linux/amd64 platform.\n",
    "2. `\"go\"` (aka. _SimpleGo_): A pure Go backend engine. Much slower than `xla`, but very portable, likely will work in all platforms supported by Go -- tested in Linux, Windows, WASM (yes, it runs in the browser).\n",
    "\n",
    "Both backends are included if you `import _ \"github.com/gomlx/gomlx/backends/default\"` like in the example below (`xla` is excluded if not Linux). \n",
    "But you can build your binary only with xla with `import _ \"github.com/gomlx/gomlx/backends/stablehlo\"` or only `go` with `import _ \"github.com/gomlx/gomlx/backends/simplego\"`.\n",
    "\n",
    "Creating a `backends.Backend` object is trivial: `backends.New()` (or `backends.MustNew()`). If the environment variable `GOMLX_BACKEND` is set, it can define the backend and its configuration (e.g.: `\"go\"`, `\"xla:cpu\"`, `\"xla:cuda\"`, etc.). Here we set it to `xla:cpu` to force the CPU backend for this tutorial.\n",
    "\n",
    "With the backend created, [`Exec`](https://pkg.go.dev/github.com/gomlx/gomlx/pkg/core/graph#Exec) is the easiest way to compile and execute computation graphs in GoMLX. \n",
    "\n",
    "To run our `Sum` function above we can do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9f99a97a-6d64-4cad-9981-43b784d569b7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Set: GOMLX_BACKEND=\"xla:cpu\"\n",
      "Short version:\t1+1=float64(2)\n",
      "Exec version:\t1+1=int64(2)\n"
     ]
    }
   ],
   "source": [
    "%env GOMLX_BACKEND=xla:cpu\n",
    "\n",
    "import _ \"github.com/gomlx/gomlx/backends/default\"\n",
    "\n",
    "// must exits with an error, logging it.\n",
    "func must[T any](value T, err error) T {\n",
    "    if err != nil {\n",
    "        log.Fatalf(\"Fatal error: %+v\", err)\n",
    "    }\n",
    "    return value\n",
    "}\n",
    "\n",
    "var backend = must(backends.New())\n",
    "\n",
    "%%\n",
    "\n",
    "// Short version: `ExecOnce` will compile the function and execute only once.\n",
    "two := must(ExecOnce(backend, /* GoMLX function */ Sum, /* Arguments */ 1.0, 1.0))\n",
    "fmt.Printf(\"Short version:\\t1+1=%v\\n\", two)\n",
    "\n",
    "// Exec version: `NewExec` will compile the Sum function so it can be executed efficiently\n",
    "// many times.\n",
    "exec := must(NewExec(backend, Sum))\n",
    "results := must(exec.Exec(1, 1))\n",
    "two = results[0]\n",
    "fmt.Printf(\"Exec version:\\t1+1=%v\\n\", two)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "329b3b41-b814-4e57-8968-bde0de4dccaa",
   "metadata": {
    "tags": []
   },
   "source": [
    "> **Note**\n",
    ">\n",
    "> - `%%` is a shortcut for `func main()`: everything after it is put inside a `main` function by **GoNB**, the Go Notebook kernel.\n",
    "> - We define the `must` function: a convenient way to handle errors for notebooks, by immediately logging the error and exiting.\n",
    "> - `backends.New()` creates a `backends.Backend` object, which connects to an accelerator if present. `backends.MustNew` is the version that panics on error. \n",
    ">   Usually one creates one at the beginning of the program and passes it around.\n",
    ">   Here **GoNB** will keep the global variable `manager` available to all cells, so we don't need\n",
    ">   to define it again.\n",
    "> - The `Exec` object created is associated with a graph building function (`Sum` in this case). It lazily \n",
    ">   compiles and executes the compiled computation as needed. Naturally the first time `Exec()` is invoked\n",
    ">   it is slow: it has to build the graph and *just-in-time* (JIT) compile it. But the compiled graph afterwards \n",
    ">   is optimized and very fast to execute, which is what we want for machine learning.\n",
    "> - The `Exec.Exec(1, 1)` method always returns a slice of results, that's why we use `[0]` to access the\n",
    ">   first result.\n",
    "> - The results of graph execution are always **tensors** (see section below). They can be converted to Go types\n",
    ">   using `.Value()` method.\n",
    "\n",
    "Something important to understand is that **Graphs have static (fixed) shapes for its inputs and outputs**. What\n",
    "it means is that, for example, if you are going to sum floats instead of ints, `Exec` would have to rebuild the graph\n",
    "to take as input two floats. Or if you want to sum a vector or matrix or ints, or any different `shapes.Shape`.\n",
    "\n",
    "> For a detailed explanation of `Shape` and the associate concepts of Axis, Dimensions and `DType` (the underlying data type), \n",
    "see [package `shapes` documentation](https://pkg.go.dev/github.com/gomlx/gomlx/pkg/core/shapes).\n",
    "\n",
    "To exemplify, let's expand our code a bit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8b1a61bf-f902-43db-a2f0-8c8ddf61375d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "- Building graph (GraphId=0) for a.shape=(Int64) and b.shape=(Int64)\n",
      "\t1+1=int64(2)\n",
      "\t0+0=int64(0)\n",
      "\t1+1=int64(2)\n",
      "\t2+2=int64(4)\n",
      "\t3+3=int64(6)\n",
      "\t4+4=int64(8)\n",
      "\n",
      "- Building graph (GraphId=1) for a.shape=(Float64) and b.shape=(Float64)\n",
      "\t3.5+1.5=float64(5)\n",
      "\n",
      "- Building graph (GraphId=2) for a.shape=(Float32)[3] and b.shape=(Float32)\n",
      "\t[1.1, 2.2, 3.3] + 10 = [3]float32{11.1, 12.2, 13.3}\n"
     ]
    }
   ],
   "source": [
    "import (\n",
    "    \"fmt\"\n",
    "    . \"github.com/gomlx/gomlx/pkg/core/graph\"\n",
    ")\n",
    "\n",
    "// Sum is the symbolic funciton (creates a graph) of summing two numbers.\n",
    "func Sum(a, b *Node) *Node {\n",
    "    g := a.Graph()  // Graph on which this symbolic computation is being built.\n",
    "    fmt.Printf(\"\\n- Building graph (GraphId=%d) for a.shape=%s and b.shape=%s\\n\", g.GraphId(), a.Shape(), b.Shape())\n",
    "    return Add(a, b)\n",
    "}\n",
    "\n",
    "\n",
    "func main() {\n",
    "    sumExec := must(NewExec(backend, Sum))\n",
    "    two := sumExec.MustExec1(1, 1)\n",
    "    fmt.Printf(\"\\t1+1=%s\\n\", two)\n",
    "\n",
    "    for ii := 0; ii < 5; ii++ {\n",
    "        sumInts := sumExec.MustExec1(ii, ii)\n",
    "        fmt.Printf(\"\\t%d+%d=%s\\n\", ii, ii, sumInts)\n",
    "    }\n",
    "\n",
    "    five := sumExec.MustExec1(3.5, 1.5)\n",
    "    fmt.Printf(\"\\t3.5+1.5=%s\\n\", five)\n",
    "\n",
    "    many := sumExec.MustExec1([]float32{1.1, 2.2, 3.3}, float32(10))\n",
    "    fmt.Printf(\"\\t[1.1, 2.2, 3.3] + 10 = %s\\n\", many)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22a1073e-bc4e-4278-a982-758bd5be4200",
   "metadata": {},
   "source": [
    "> **Note**\n",
    ">   - Each time new input shapes are found, a new graph is created. We added a `fmt.Printf` to tell us\n",
    ">     the GraphId and more importantly the shape of the graph operands.\n",
    ">     Notice that `fmt.Printf` is not part of the computation graph, it's only part of the graph building function.\n",
    ">     We'll see later `Node.SetLogged` on how to print intermediary results in the middle of the execution of the graph.\n",
    ">   - Every `Node` has an associated shape (`shapes.Shape` type). A shape is defined by its underlying data type \n",
    ">     `shapes.DType` and its axes dimensions. For scalars, the shape has zero axes (dimensions). E.g.: `(Int64)[]` represents\n",
    ">     a scalar `int` value, and `(Float32)[3]` represents a vector with 3 `float32` values. More details\n",
    ">     and the list of data types (aka. _dtype_) supported in the package [github.com/gomlx/gopjrt/dtypes](https://pkg.go.dev/github.com/gomlx/gopjrt/dtypes).\n",
    ">   - `Exec.MustExec1` is an alias to the previous `Exec.Exec1` that conveniently: (a) panics on error; (b) returns the first result.\n",
    "\n",
    "In general the graph operations only work with the same `dtypes.DType`. \n",
    "If they are different, they are reported back with a `panic` (works like an exception, and can be caught) with an error with a full stack-trace in the returned result.\n",
    "And with that, let's talk how error is handled in GoMLX computational graph building:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93721733-8372-41dd-a9f8-e822913be569",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### `Exec` aliases:\n",
    "\n",
    "Because the `graph.Exec` (and later the similar `context.Exec`, in a few sections below) is so central, we provide some aliases to make it more \"ergonomic\" to use.\n",
    "\n",
    "To create an `Exec` object:\n",
    "\n",
    "- `NewExec`, `MustNewExec`: the first returns an error, the second panics.\n",
    "\n",
    "To create, execute and finalize an `Exec` object in one line, for when we want to execute a computation only once:\n",
    "\n",
    "- `Exec.ExecOnce` and `Exec.ExecOnceN`: the first can be used if the computation returns only one result. The second if the computation returns > 1 result, it returns a slice of tensors (see below).\n",
    "\n",
    "To execute an `Exec` object:\n",
    "\n",
    "- `Exec.Exec`, `Exec.Exec1`, `Exec.Exec2`, ... `Exec.Exec4`: The first returns a slice of tensors; `Exec1` returns only one tensor, it assumes\n",
    "  the computation graph has only one output; `Exec2` returns 2 tensors, and so on. They all also return an error.\n",
    "- `Exec.ExecWithGraph` executes the computation and returns the results (as a slice of tensors) as well as the `graph.Graph` object that holds the compiled computation.\n",
    "\n",
    "There is always a `MustExec...` version of each method that panics on error, instead of returning them."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dc80803-6db2-4672-9197-f4a045d21628",
   "metadata": {},
   "source": [
    "## Error Handling\n",
    "\n",
    "During the _graph building_ (or what we call symbolic computation) **GoMLX** diverges from the usual Golang idiomatic error handling: it uses exceptions in the form of a `panic()` (with an `error` object with a stack trace).\n",
    "\n",
    "Usually, the symbolic computation functions (like the `Sum` function we defined above) are called by the `Exec.Exec` methods. And they capture any panics and return them as errors to the caller.\n",
    "\n",
    "> **Note**:\n",
    "> The author is aware this is controvertial for some. So below are the motivations to use exceptions, as opposed to\n",
    "> returning an error everywhere:\n",
    "> - While implementing long mathematical functions, introducing error handling in between each operation would\n",
    ">   be too distracting, and unwieldly. Notice the Go language makes the same choice; e.g.: that's why the division operator `/`\n",
    ">   doesn't return the result of the division and an error, and instead panics with division by zero.\n",
    "> - Speed of building the computation graph is not a concern: it has no impact on the execution of the\n",
    ">   JIT-compiled code later on.\n",
    "\n",
    "Let's create an example with an error to see how this works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dce20082-e817-47a4-ac3e-9eec27a6e3a5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "- Building graph (GraphId=0) for a.shape=(Float64) and b.shape=(Int64)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025/12/05 10:56:45 Fatal error: cannot broadcast Float64 and Int64 for \"Add\": they have different dtypes\n",
      "github.com/gomlx/gomlx/backends/stablehlo.(*Builder).broadcastForBinaryOps\n",
      "\t/home/janpf/Projects/gomlx/backends/stablehlo/builder.go:347\n",
      "github.com/gomlx/gomlx/backends/stablehlo.(*Builder).Add\n",
      "\t/home/janpf/Projects/gomlx/backends/stablehlo/gen_binary_ops.go:13\n",
      "github.com/gomlx/gomlx/pkg/core/graph.Add\n",
      "\t/home/janpf/Projects/gomlx/pkg/core/graph/gen_backend_ops.go:194\n",
      "main.Sum\n",
      "\t \u001b[7m[[ Cell Line 0 ]]\u001b[0m /tmp/gonb_9dbcd313/main.go:20\n",
      "reflect.Value.call\n",
      "\t/snap/go/current/src/reflect/value.go:581\n",
      "reflect.Value.Call\n",
      "\t/snap/go/current/src/reflect/value.go:365\n",
      "github.com/gomlx/gomlx/pkg/core/graph.(*Exec).createAndCacheGraph\n",
      "\t/home/janpf/Projects/gomlx/pkg/core/graph/exec.go:888\n",
      "github.com/gomlx/gomlx/pkg/core/graph.(*Exec).findOrCreateGraph\n",
      "\t/home/janpf/Projects/gomlx/pkg/core/graph/exec.go:980\n",
      "github.com/gomlx/gomlx/pkg/core/graph.(*Exec).compileAndExecute.func1\n",
      "\t/home/janpf/Projects/gomlx/pkg/core/graph/exec.go:704\n",
      "github.com/gomlx/gomlx/internal/exceptions.TryCatch[...]\n",
      "\t/home/janpf/Projects/gomlx/internal/exceptions/exceptions.go:89\n",
      "github.com/gomlx/gomlx/pkg/core/graph.(*Exec).compileAndExecute\n",
      "\t/home/janpf/Projects/gomlx/pkg/core/graph/exec.go:703\n",
      "github.com/gomlx/gomlx/pkg/core/graph.(*Exec).ExecWithGraphOnDevice\n",
      "\t/home/janpf/Projects/gomlx/pkg/core/graph/exec.go:498\n",
      "github.com/gomlx/gomlx/pkg/core/graph.(*Exec).ExecWithGraph\n",
      "\t/home/janpf/Projects/gomlx/pkg/core/graph/execaliases.go:105\n",
      "github.com/gomlx/gomlx/pkg/core/graph.(*Exec).Exec\n",
      "\t/home/janpf/Projects/gomlx/pkg/core/graph/execaliases.go:35\n",
      "main.main\n",
      "\t \u001b[7m[[ Cell [5] Line 3 ]]\u001b[0m /tmp/gonb_9dbcd313/main.go:35\n",
      "runtime.main\n",
      "\t/snap/go/current/src/runtime/proc.go:285\n",
      "runtime.goexit\n",
      "\t/snap/go/current/src/runtime/asm_amd64.s:1693\n",
      "exit status 1\n"
     ]
    }
   ],
   "source": [
    "%%\n",
    "sumExec := MustNewExec(backend, Sum)\n",
    "_ = must(sumExec.Exec(1.1, 2)) // Error: arguments have different dtypes Float64 and Int64."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1622fe6c-cbc8-46fd-8d56-eaec55e3fc9d",
   "metadata": {},
   "source": [
    "> **Note**\n",
    "> - In the stack-trace above there are 2 lines of interest, that typically help to debug such issues:\n",
    ">   1. Where in the graph building function `main.Sum` function the invalid operation was created: **Line 9 of the previous cell**.\n",
    ">   2. Where in the `main` function, the graph was attempted to be executed: **Line 3 of this cell**.\n",
    "> - You can enable displaying line-numbers in the JupyterLab with \"ESC+L\" (upper-case L).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52564b67-ae23-4545-b853-9a5c1fe384e2",
   "metadata": {},
   "source": [
    "## Tensors\n",
    "\n",
    "Tensors are multidimensional arrays of a given data type (`dtypes.DType`) defined in the package [`github.com/gomlx/gomlx/pkg/core/tensors`](https://pkg.go.dev/github.com/gomlx/gomlx/pkg/core/tensors). \n",
    "\n",
    "For GoMLX tensors work as **containers of data** that are used as concrete inputs and outputs for the\n",
    "execution of computational graphs. There is only basic support to manipulate tensors directly (it includes direct access to its data): instead usually one operates on the tensor using computational graphs and the `Exec` object. Tensors have a shape (`shapes.Shape`) just like a `graph.Node`.\n",
    "\n",
    "When talking about a tensor, we are referring to the package's `tensors.Tensor` object (a pointer to it).\n",
    "\n",
    "The tensor raw values can be stored either locally (as a Go slice) or on the device where the graph is being executed. E.g.: a GPU, or even a CPU, in which case it is owned by the `backends.Backend`.\n",
    "\n",
    "While the `tensors.Tensor` object handles the synchronization automatically among the local/on-device versions, it's important\n",
    "that you are aware that there is this distinction. There is a time cost to transfer data from, for example, a GPU to the CPU and back.\n",
    "\n",
    "Effectively, what that means is that, during training, you dont' want to transfer anything locally often,\n",
    "for instance to save a checkpoint. Most of this is handled automatically by **GoMLX** libraries, but not everything. \n",
    "Example: because some of the resources are allocated in the accelerator (GPU) and not managed by Go, the garbage collector may not be aware of the memory pressure on these devices--there is a `Tensor.Finalize` call that immediately releases the data associated with a tensor.\n",
    "\n",
    "Whenever a graph is executed (with `Exec.Exec`), non-tensor input values are automatically converted to tensors, and those are transferred to the device doing the execution automatically.\n",
    "\n",
    "Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6acb053e-13e3-4d79-8cc5-5cb134d0e78a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "counter.type=*tensors.Tensor, counter=int64(1)\n",
      "counter=int64(11)\n"
     ]
    }
   ],
   "source": [
    "%%\n",
    "onePlusExec := MustNewExec(backend, func (x *Node) *Node {\n",
    "    return OnePlus(x)\n",
    "})\n",
    "// exec.MustExec1 will return a *tensor.Tensor.\n",
    "counter := onePlusExec.MustExec1(0)\n",
    "// counter.String() will first transfer counter to local (using counter.Local()) to print its values.\n",
    "fmt.Printf(\"counter.type=%T, counter=%s\\n\", counter, counter.String())\n",
    "for ii := 0; ii < 10; ii++ {\n",
    "    // Since the counter is not being used locally between the calls, the tensor will only use the\n",
    "    // device storage.\n",
    "    counter = onePlusExec.MustExec1(counter)\n",
    "}\n",
    "// Again counter.String() (called implicitly by fmt.Printf) will first transfer the counter value locally, and then convert to a Go value.\n",
    "fmt.Printf(\"counter=%v\\n\", counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ea8aeee-e88d-43c5-8a91-362099fce822",
   "metadata": {},
   "source": [
    "> **Note**:\n",
    "> - In the first call to `onePlusExec.MustExec1(0)`, the Go constant `0` is automatically converted to a `*tensor.Tensor`\n",
    ">   by the `graph.Exec` and fed to the graph. It returns a `[]*tensor.Tensor` with one element, containing `0+1=1`.\n",
    "> - The returned tensor is initially only stored on device. But when we print it, it is automatically transferred to Go.\n",
    "> - While executing the loop the counter will not be transferred locally, everything happens efficiently in the accelerator device (GPU, TPU, or even the CPU executor).\n",
    "> - When we print the final result in `counter.String()` again the data is transferred locally automatically."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98f8f8c2-9716-4212-87ef-7dc7b884c11f",
   "metadata": {},
   "source": [
    "### New tensors\n",
    "\n",
    "There are several ways to create tensors, the most common:\n",
    "\n",
    " - `tensors.FromValue[S](value S)`: Generics conversion, works with any supported `DType` scalar\n",
    "   as well as with any arbitrary multidimensional slice. Slices of rank > 1 must be regular, that is\n",
    "   all the sub-slices must have the same shape. E.g: `FromValue([][]float{{1,2}, {3, 5}, {7, 11}})`.\n",
    "   There is also a non-generic version `tensors.FromAnyValue(value any)`. This is what `Exec.Exec()`\n",
    "   uses to convert arbitrary values to a tensor.\n",
    " - `FromShape(shape shapes.Shape)`: creates a Local tensor with the given shape, and initialized to zero. See\n",
    "   documentation on `Tensor.MutableFlatData(...)` to mutate tensors in place.\n",
    " - `FromScalarAndDimensions[T](value T, dimensions ...int)`: creates a tensor with the\n",
    "   given dimensions, filled with the scalar value given. `T` must be one of the supported types.\n",
    "\n",
    "You can print the contents of a `tensors.Tensor` simply using `Tensor.String()`, or convert it back to the\n",
    "corresponding Go type using `Tensor.Value()` (it returns an `any` that you can cast to the expected type).\n",
    "\n",
    "See more documentation in [github.com/gomlx/gomlx/pkg/core/tensors](https://pkg.go.dev/github.com/gomlx/gomlx/pkg/core/tensors).\n",
    "\n",
    "Errors in the manipulation of Tensors (e.g. invalid values) are also reported back with exception thrown with `panic`, with full stack-traces -- they only happen due to invalid inputs (nil values, wrong data types, etc.). The errors can easily be caught (with `recover()` or with `exceptions.TryCatch` helper) when needed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd1593dc-b621-4307-b9fb-06376ccf1940",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Gradients\n",
    "\n",
    "An important functionality required to train machine learning models based on gradient descent is calculating the gradients of some value (loss) being optimized with respect to some variable / quantity. \n",
    "\n",
    "**GoMLX** does this symbolically, during graph building time. It adds to the graph the computation for the gradient.\n",
    "\n",
    "Example: let's calculate the gradient of the function $f(x, y) = x^2 + xy$ for a few values of $x$ and $y$.\n",
    "Algebraically we have:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "f(x, y) &= x^2 + xy \\\\\n",
    "df/dx(x,y) &= 2x + y \\\\\n",
    "df/dy(x,y) &= x \\\\\n",
    "\\end{aligned}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a1736ccc-df7c-4b8f-a6c2-312719547f69",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "f(x=[0 1 2], y=[10 20 30])=[3]float64{0, 21, 64},\n",
      "\tdf/dx=[3]float64{10, 22, 34},\n",
      "\tdf/dy=[3]float64{0, 1, 2}\n"
     ]
    }
   ],
   "source": [
    "func f(x, y *Node) *Node {\n",
    "    return Add(Square(x), Mul(x, y))\n",
    "}\n",
    "\n",
    "func gradOfF(x, y *Node) (output, gradX, gradY *Node) {\n",
    "    output = f(x, y)\n",
    "    reduced := ReduceAllSum(output) // In case x and y are not scalars.\n",
    "    grads := Gradient(reduced, x, y)\n",
    "    gradX, gradY = grads[0], grads[1] // df/dx, df/dy\n",
    "    return output, gradX, gradY\n",
    "}\n",
    "\n",
    "%%\n",
    "exec := MustNewExec(backend, gradOfF)\n",
    "x := []float64{0, 1, 2}\n",
    "y := []float64{10, 20, 30}\n",
    "output, gradX, gradY := exec.MustExec3(x, y)\n",
    "fmt.Printf(\"f(x=%v, y=%v)=%s,\\n\\tdf/dx=%s,\\n\\tdf/dy=%s\\n\", x, y, output, gradX, gradY)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4afbd1b1-d53f-45be-8b95-8aab31bd3de2",
   "metadata": {},
   "source": [
    "> **Note**:\n",
    "> - As of yet, GoMLX only calculates gradients of a scalar (typically a model loss) with respect to arbitrary tensors.\n",
    ">   It does not yet calculate **jacobians**, that is, if the value we are deriving is not a scalar.\n",
    "> - A question that may arise is whether it calculates the second derivative (*hessian*).\n",
    ">   In principle the machinery to do that is in place, but there are 2 limitations:\n",
    ">   (1) not all operations have their derivative implemented, in particular some of the operations that are only used when calculating the first derivative;\n",
    ">   (2) it only calculates the gradient with respect to a scalar, in most cases the hessian will be the gradient of a gradient,\n",
    ">   usually of higher rank.\n",
    "> - It's not impossible (or hard) to add support for those. If you need it, contributions to the project here are welcome ;)  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "83861fda-d676-48a0-b5ba-bdbaa93cb2fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ". removed func gradOfF\n",
      ". removed func f\n"
     ]
    }
   ],
   "source": [
    "// Removing the previous definitions of `f` and `gradOfF`\n",
    "%rm gradOfF f"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b532f529-5e83-4aed-a70d-e163594fafec",
   "metadata": {},
   "source": [
    "# Machine Learning with GoMLX\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1bc04a1-30df-4a5a-aede-07dedf2a3798",
   "metadata": {},
   "source": [
    "## Variables and Context\n",
    "\n",
    "Computation graphs are [pure functions](https://en.wikipedia.org/wiki/Pure_function): they have no state,\n",
    "they take inputs, return outputs and everything in between is transient.\n",
    "\n",
    "For Machine Learning as well as many types of computations, it's convenient to use and update \"contextual\" information\n",
    "about a model being trained or executed, like its weights/parameters.\n",
    "We want to be able to easily access and update these weights in the middle of a computation graph.\n",
    "\n",
    "To allow this **GoMLX** offers the [`context.Context`](https://pkg.go.dev/github.com/gomlx/gomlx/pkg/ml/context) object \n",
    "(completely unrelated to the usual Go's `context` package), and a corresponding `context.Exec`. \n",
    "It is a container of **variables**, and it manages automatically its updates, passing it as \"side\n",
    "inputs\" and returning them as \"side outputs\" if they were updated in the computation graph.\n",
    "\n",
    "This may sound more complex than it is in practice. Let's see a simple example of maintaining a counter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "92acdfcb-bb9a-4665-a879-a6b4596408a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counting:\n",
      "\tcount=int32(11)\n",
      "\tcount=int32(12)\n",
      "\tcount=int32(13)\n",
      "- State of counter=%!s(int32=13)\n"
     ]
    }
   ],
   "source": [
    "import \"github.com/gomlx/gomlx/pkg/ml/context\"\n",
    "\n",
    "%%\n",
    "ctx := context.New()\n",
    "exec := context.MustNewExec(backend, ctx, func(ctx *context.Context, g *Graph) *Node {\n",
    "    counterVar := ctx.VariableWithValue(\"count\", int32(10))\n",
    "    count := counterVar.ValueGraph(g)\n",
    "    count = AddScalar(count, 1)\n",
    "    counterVar.SetValueGraph(count)\n",
    "    return count\n",
    "})\n",
    "fmt.Println(\"Counting:\")\n",
    "for range(3) {\n",
    "    fmt.Printf(\"\\tcount=%s\\n\", exec.MustExec1())\n",
    "}\n",
    "counterVar := must(ctx.InspectVariable(ctx.Scope(), \"count\").Value())\n",
    "fmt.Printf(\"- State of counter=%s\\n\", counterVar.Value())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "804ecc5b-cdb2-4177-acc6-e839c991b2ef",
   "metadata": {},
   "source": [
    "> **Note**:\n",
    "> - We are using `context.Exec`, while before we were using `graph.Exec`. The main difference is that\n",
    ">   `context.Exec` compiles and executes graph functions that take a context as its first parameter, and it \n",
    ">   automatically handles the passing of variables as side inputs and outputs (for those variables updated)\n",
    ">   of the computation graph.\n",
    "> - During graph building, we access and set the variables's symbolic values (`graph.Node`) with `Variable.ValueGraph` and `VariableSetValueGraph`.\n",
    "> - Outside graph building, we can access the last value set to a variable by using `Variable.Value()` and \n",
    ">   `Variable.SetValue`. They return/take concrete `*tensor.Tensor` types.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "955465cb-f563-43d7-90e2-d23e6b762c3d",
   "metadata": {},
   "source": [
    "### Finding $argmin_{x}{f(x)}$ example\n",
    "\n",
    "A more elaborate example: let's find $argmin_{x}{f(x)}$ where $f(x) = ax^2 + bx + c$. \n",
    "\n",
    "If we solve it literally we should get, for $a > 0$, $argmin_{x}{f(x)} = \\frac{-b}{2a}$. \n",
    "\n",
    "Instead, we solve it numerically, using gradient descent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "610daa58-952e-4d8e-8769-c2d8609baa14",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import (\n",
    "    \"flag\"\n",
    "    \"github.com/gomlx/gomlx/pkg/core/shapes\"\n",
    ")\n",
    "\n",
    "var (\n",
    "    flagA = flag.Float64(\"a\", 1.0, \"Value of a in the equation ax^2+bx+c\")\n",
    "    flagB = flag.Float64(\"b\", 2.0, \"Value of b in the equation ax^2+bx+c\")\n",
    "    flagC = flag.Float64(\"c\", 4.0, \"Value of c in the equation ax^2+bx+c\")\n",
    "    flagNumSteps = flag.Int(\"steps\", 10, \"Number of gradient descent steps to perform\")\n",
    "    flagLearningRate    = flag.Float64(\"lr\", 0.1, \"Initial learning rate.\")\n",
    ")\n",
    "\n",
    "// f(x) = ax^2 + bx + c\n",
    "func f(x *Node) *Node {\n",
    "    f := MulScalar(Square(x), *flagA)\n",
    "    f = Add(f, MulScalar(x, *flagB))\n",
    "    f = AddScalar(f, *flagC)\n",
    "    return f\n",
    "}\n",
    "\n",
    "// minimizeF does one gradient descent step on F by moving a variable \"x\",\n",
    "// and returns the value of the function at the current \"x\".\n",
    "func minimizeF(ctx *context.Context, graph *Graph) *Node {\n",
    "    // Create or reuse existing variable \"x\" -- no graph operation is created with this, it's\n",
    "    // only a reference.\n",
    "    xVar := ctx.VariableWithShape(\"x\", shapes.Make(dtypes.Float64))\n",
    "\n",
    "    x := xVar.ValueGraph(graph)                        // Read variable for the current graph.\n",
    "    y := f(x)                                          // Value of f(x).\n",
    "\n",
    "    // Gradient always return a slice, we take the first element for grad of X.\n",
    "    gradX := Gradient(y, x)[0]\n",
    "\n",
    "    // stepNum += 1\n",
    "    stepNumVar := ctx.VariableWithValue(\"stepNum\", 0.0)  // Creates the variable if not existing, or retrieve it if already exists.\n",
    "    stepNum := stepNumVar.ValueGraph(graph)\n",
    "    stepNum = OnePlus(stepNum)\n",
    "    stepNumVar.SetValueGraph(stepNum)\n",
    "\n",
    "    // step = -learningRate * gradX / Sqrt(stepNum)\n",
    "    step := Div(gradX, Sqrt(stepNum))\n",
    "    step = MulScalar(step, -*flagLearningRate)\n",
    "\n",
    "    // x += step\n",
    "    x = Add(x, step)\n",
    "    xVar.SetValueGraph(x)\n",
    "    return y  // f(x)\n",
    "}\n",
    "\n",
    "func Solve() {\n",
    "    ctx := context.New()\n",
    "    exec := context.MustNewExec(backend, ctx, minimizeF)\n",
    "    for ii := 0; ii < *flagNumSteps-1; ii++ {\n",
    "        _ = exec.MustExec()\n",
    "    }\n",
    "    y := exec.MustExec1()\n",
    "    x := must(ctx.InspectVariable(ctx.Scope(), \"x\").Value())\n",
    "    stepNum := must(ctx.InspectVariable(ctx.Scope(), \"stepNum\").Value())\n",
    "    fmt.Printf(\"Minimum found at x=%g, f(x)=%g after %f steps.\\n\", x.Value(), y.Value(), stepNum.Value())\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14d308a2-582b-49ce-b1e4-d55997b66f97",
   "metadata": {},
   "source": [
    "The code above created `Solve()` that will solve for the values set by the flags `a`, `b`, and `c`.\n",
    "\n",
    "Let's try a few values:\n",
    "\n",
    "> **Note**: `%%` in GoNB automatically creates a `func main()` and passes the extra arguments to the Go program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f7ba4841-21a4-4715-ac9e-6dba59b97218",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Minimum found at x=-1, f(x)=2 after 10.000000 steps.\n"
     ]
    }
   ],
   "source": [
    "%% --a=1 --b=2 --c=3 --steps=10 --lr=0.5\n",
    "Solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9ff74708-1bed-49ce-b3c8-b02e9278d8fc",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Minimum found at x=-3, f(x)=2 after 10.000000 steps.\n"
     ]
    }
   ],
   "source": [
    "%% --a=2 --b=12 --c=20 --steps=10 --lr=0.5\n",
    "Solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "082674a7-937d-48e0-8f29-77c5bb3eee83",
   "metadata": {},
   "source": [
    "> **Note**:\n",
    "> - We created two variables, one for \"x\" that we were optimizing to minimize $f(x)$, and one variable \"stepNum\",\n",
    ">   used to keep track how many steps were already executed. \n",
    "> - Yes, if we set `--lr=1` (the learning rate), it will get to the minimum in one step for the quadratic f(x). 😉\n",
    "\n",
    "There is more to [`context.Context`](https://pkg.go.dev/github.com/gomlx/gomlx/pkg/ml/context), some we'll present \n",
    "on the next section on *Machine Learning*, others can be found in its documentation. A few things worth advancing:\n",
    "\n",
    "* `Context` is always configured at a certain *scope*, and variables are unique within its scope. Scope\n",
    "  is easily changed with `ctx.In(\"new_scope\")`. So the `Context` object is a scope (a string) and a pointer to \n",
    "  the actual data (variables, graph and model parameters).\n",
    "* `Context` also holds model parameters (concrete Go values), which are also scoped. \n",
    "  Those can be hyperparameters for the models (learning rate, regularization, etc.) or anything the user or\n",
    "  any library may create a convention for. They are more convenient than using hundreds of flags.\n",
    "  See example of its usage in the [UCI Adult](https://github.com/gomlx/gomlx/blob/main/examples/adult/demo/main.go) example.\n",
    "* Similarly `Context` also holds \"Graph parameters\". Those are very similar to model parameters, but they\n",
    "  have one value per Graph. So if a model is created with parameters of different shape (or for training/evaluation),\n",
    "  each version will have its own Graph parameters. Don't worry about this now -- if you need it later\n",
    "  when building complex graphs, the funcitonality will be there."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee9edc1e-c79c-4a8e-b570-b20c0195db44",
   "metadata": {},
   "source": [
    "## Training a machine learning model\n",
    "\n",
    "The previous sections presented the fundamentals of what is needed to implement machine learning. This section\n",
    "we present various sub-packages that provide high level ML layers and tools that make building, training and \n",
    "evaluating a model trivial.\n",
    "\n",
    "First is the package `layers` (see code in [ml/layers](../../../tree/main/ml/layers/). It provides several composable ML layers. \n",
    "These are graph building functions, most of which take a `*context.Context` as first parameter, where they store \n",
    "variables or access hyperparameters. \n",
    "There are several such layers, for example: `layers.Dense`, `layers/fnn.New`, `layers/kan.New`,  \n",
    "`layers.Dropout`, `layers.PiecewiseLinearCalibration` (very good for normalization of inputs), `layers.BatchNorm`, \n",
    "`layers.LayerNorm`, `layers.Convolution`, `layers.MultiHeadAttention` (for [Transformers](https://arxiv.org/abs/1706.03762) \n",
    "layers), etc.\n",
    "\n",
    "The package `train` offers two main functionalities: `train.Trainer` will build a *train step* and an *eval step*\n",
    "graph, given a model graph building function and an optimizer. This graph can be executed in sequence to train a model.\n",
    "The package also provides `train.Loop` that simply loop over a `train.Dataset` interface, reading data and feeding\n",
    "it to `Trainer.TrainStep` or `Trainer.EvalStep`, along with executing configurable hooks on the training loop. One\n",
    "such hooks is provided by `gomlx/train/commandline.AttachProgressBar(loop)`, it pretty prints the progress during\n",
    "training on the command line.\n",
    "\n",
    "There are also a collection of **optimizers**, **initializers**, **loss functions**, **metrics**, etc. \n",
    "For any functionality there is always an example under the [examples/](../../../tree/main/examples/) subdirectory.\n",
    "\n",
    "Let's look at the simplest ML example: [`linear`](../../../tree/main/examples/linear/linear.go), \n",
    "which trains a linear model on nosiy generated data.\n",
    "\n",
    "Here are the constants of our problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d82352a6-2c89-4c00-b483-522f7c067f25",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "const (\n",
    "    CoefficientMu    = 0.0\n",
    "    CoefficientSigma = 5.0\n",
    "    BiasMu           = 1.0\n",
    "    BiasSigma        = 10.0\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a176818-1ad3-416c-8de9-c52bc3a83c97",
   "metadata": {},
   "source": [
    "To generate synthetic data it first randomly chose some random coefficients and bias based on which data is \n",
    "generated. These selected coefficients is the ones we want to try to learn using ML. The coefficients could\n",
    "have been selected in Go directly using `math/random`, but just for fun, we do it using a computation graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "83231803-576d-4504-9978-b32175ee7cc6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Example of target: coefficients=[0.722 -7.14 -5.2], bias=-5.19\n"
     ]
    }
   ],
   "source": [
    "import (\n",
    "    \"github.com/gomlx/gomlx/pkg/core/shapes\"\n",
    "    \"github.com/gomlx/gomlx/pkg/core/tensors\"\n",
    ")\n",
    "\n",
    "// initCoefficients chooses random coefficients and bias. These are the true values the model will\n",
    "// attempt to learn.\n",
    "func initCoefficients(backend backends.Backend, numVariables int) (coefficients, bias *tensors.Tensor) {\n",
    "    e := MustNewExec(backend, func(g *Graph) (coefficients, bias *Node) {\n",
    "        rngState := Const(g, must(RNGState()))\n",
    "        rngState, coefficients = RandomNormal(rngState, shapes.Make(dtypes.Float64, numVariables))\n",
    "        coefficients = AddScalar(MulScalar(coefficients, CoefficientSigma), CoefficientMu)\n",
    "        rngState, bias = RandomNormal(rngState, shapes.Make(dtypes.Float64))\n",
    "        bias = AddScalar(MulScalar(bias, BiasSigma), BiasMu)\n",
    "        return\n",
    "    })\n",
    "    coefficients, bias = e.MustExec2()\n",
    "    return\n",
    "}\n",
    "\n",
    "%%\n",
    "coef, bias := initCoefficients(backend, 3)\n",
    "fmt.Printf(\"Example of target: coefficients=%0.3v, bias=%0.3v\\n\", coef.Value(), bias.Value())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92557ffc-7c1a-414d-9efe-f9eb0e2f4a07",
   "metadata": {
    "tags": []
   },
   "source": [
    "> **Note**\n",
    "> - This code should look familiar, using things we presented earlier in the tutorial. It creates a computation\n",
    ">   graph to generate randomly the `coefficients` and `bias`. Then it executes it and returns the result. \n",
    "> - Notice that since the *computation graph is functional*: we need to pass around the random number\n",
    ">   generator state, which gets updated at each call to `RandomUniform` or `RandomNormal`.\n",
    ">   - Alternatively the `context.Context` introduced earlier can keep the state\n",
    ">     as a variable, and provides a simpler interface: see `Context.RandomUniform` and `Context.RandomNormal`.\n",
    "\n",
    "Next, we want to generate the data (examples): we generate random inputs, and then the label using the\n",
    "selected coefficients plus some normal noise.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1ee57593-a9dd-4038-a704-694a120840fa",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Target: coefficients=[-3.22 -5.42 7.52], bias=-4.46\n",
      "5 dataset examples:\n",
      "\tx=[-0.139 0.667 -1.94]; label=[-22.2]\n",
      "\tx=[-1.26 1.19 0.0864]; label=[-5.85]\n",
      "\tx=[-0.483 -0.322 0.102]; label=[-0.456]\n",
      "\tx=[0.953 -0.899 -0.292]; label=[-4.54]\n",
      "\tx=[0.425 0.25 -2.06]; label=[-22.8]\n"
     ]
    }
   ],
   "source": [
    "func buildExamples(backend backends.Backend, coef, bias *tensors.Tensor, numExamples int, noise float64) (inputs, labels *tensors.Tensor) {\n",
    "    e := MustNewExec(backend, func(coef, bias *Node) (inputs, labels *Node) {\n",
    "        g := coef.Graph()\n",
    "        numFeatures := coef.Shape().Dimensions[0]\n",
    "\n",
    "        // Random inputs (observations).\n",
    "        rngState := Const(g, must(RNGState()))\n",
    "        rngState, inputs = RandomNormal(rngState, shapes.Make(coef.DType(), numExamples, numFeatures))\n",
    "        coef = ExpandDims(coef, 0)\n",
    "\n",
    "        // Calculate perfect labels.\n",
    "        labels = ReduceAndKeep(Mul(inputs, coef), ReduceSum, -1)\n",
    "        labels = Add(labels, bias)\n",
    "        if noise > 0 {\n",
    "            // Add some noise to the labels.\n",
    "            var noiseVector *Node\n",
    "            rngState, noiseVector = RandomNormal(rngState, labels.Shape())\n",
    "            noiseVector = MulScalar(noiseVector, noise)\n",
    "            labels = Add(labels, noiseVector)\n",
    "        }\n",
    "        return\n",
    "    })\n",
    "    inputs, labels = e.MustExec2(coef, bias)\n",
    "    return\n",
    "}\n",
    "\n",
    "%%\n",
    "coef, bias := initCoefficients(backend, 3)\n",
    "numExamples := 5\n",
    "inputsTensor, labelsTensor := buildExamples(backend, coef, bias, numExamples, 0.2)\n",
    "fmt.Printf(\"Target: coefficients=%0.3v, bias=%0.3v\\n\", coef.Value(), bias.Value())\n",
    "\n",
    "fmt.Printf(\"%d dataset examples:\\n\", numExamples)\n",
    "inputs := inputsTensor.Value().([][]float64)\n",
    "labels := labelsTensor.Value().([][]float64)\n",
    "for ii := 0; ii < numExamples; ii ++ {\n",
    "    fmt.Printf(\"\\tx=%0.3v; label=%0.3v\\n\", inputs[ii], labels[ii])\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c2ef923-f7cb-4011-ab84-901e0747cd68",
   "metadata": {},
   "source": [
    "### Dataset\n",
    "\n",
    "Now the first new concept of this section: `train.Dataset` is the interface that is used to feed data to\n",
    "during a training loop or evaluation. \n",
    "\n",
    "There are three methods: `Dataset.Yield` that returns the next batch of examples; \n",
    "`Dataset.Reset` restarts the dataset, for datasets that don't loop indefinitely; \n",
    "Finally `Dataset.Name` returns the dataset name, usually used for metric names, logging and printing.\n",
    "\n",
    "Datasets also yield a `spec`, an opaque type for GoMLX (defined as `any`), that allows the dataset to communicate\n",
    "to the model which type of data it is generating. In our case, since it's always the same data, we don't need it, \n",
    "so we keep it set to `nil`. \n",
    "If one would implement a dataset like a generic CSV file, one may want to communicate \n",
    "to the model the field names to the Model throught the `spec`, for instance. See the documentation for more details.\n",
    "\n",
    "For our linear synthetic data we implement the simplest `train.Dataset`: the whole data is pre-generated, and we return a giant batch with the full data every time -- we could also have more easily used `data.InMemoryDataset` for that, but for didactic purposes we create it from scratch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "2cadb18d-ef0e-4036-b450-a64f8286a65b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import \"github.com/gomlx/gomlx/pkg/ml/train\"\n",
    "\n",
    "// TrivialDataset always returns the whole data.\n",
    "type TrivialDataset struct {\n",
    "    name string\n",
    "    inputs, labels []*tensors.Tensor\n",
    "}\n",
    "\n",
    "var (\n",
    "    // Assert Dataset implements train.Dataset.\n",
    "    _ train.Dataset = &TrivialDataset{}\n",
    ")\n",
    "// Name implements train.Dataset.\n",
    "func (ds *TrivialDataset) Name() string { return ds.name }\n",
    "\n",
    "// Yield implements train.Dataset.\n",
    "func (ds *TrivialDataset) Yield() (spec any, inputs, labels []*tensors.Tensor, err error) {\n",
    "    return ds, ds.inputs, ds.labels, nil\n",
    "}\n",
    "\n",
    "// IsOwnershipTransferred tells the training loop that the dataset keeps ownership of the yielded tensors.\n",
    "func (ds *TrivialDataset) IsOwnershipTransferred() bool {\n",
    "    return false\n",
    "}\n",
    "\n",
    "// Reset implements train.Dataset.\n",
    "func (ds *TrivialDataset) Reset() {}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1647e85-76c8-4adc-b317-2c50356231c4",
   "metadata": {},
   "source": [
    "> **Note**:\n",
    "> * More often it is more work pre-processing data than actually building an ML model ... that's life 🙁\n",
    "> * In the [examples/](../../../tree/main/examples/) subdirectory we implement `train.Dataset` for some\n",
    ">   well known data sets: UCI Adult, Cifar-10, Cifar-100, IMDB Reviews, Kaggle's Dogs vs Cats, Oxford Flowers 102.\n",
    ">   These can be used as libraries to easily try different models. \n",
    ">   If you are workig on public datasets, please consider contributing similar libraries.\n",
    "> * **GoMLX** also include the `data.InMemoryDataset`, which can be created from tensors in one line \n",
    "\n",
    "The package [github.com/gomlx/gomlx/pkg/ml/datasets](https://pkg.go.dev/github.com/gomlx/gomlx/pkg/ml/data) provides several\n",
    "tools to facilitate the work here:\n",
    "\n",
    "* `Parallel`: parallelizes any dataset, includes some buffer.\n",
    "* `InMemory`: reads a dataset into (accelerator) memory, and then serves it from there -- greatly accelerates training.\n",
    "  We could have used that instead of defining `TrivialDataset`, but we left it because it's common to specialize\n",
    "  datasets.\n",
    "* `ConstantDataset`: trivial Dataset that returns always empty inputs and labels. Used when one is going to generate\n",
    "  the data dynamically -- it could also have been used for this example.\n",
    "  See example in [KANs Shapes demo](https://github.com/gomlx/gomlx/blob/main/examples/discretekan/kans_shapes.ipynb).\n",
    "* Downloading of datasets (with progress-bar) and checksum functions.\n",
    "\n",
    "### ModelFn\n",
    "\n",
    "Next we build a model, that for our `train` package means implementing a function with the following signature:\n",
    "\n",
    "```go\n",
    "type ModelFn func(ctx *context.Context, spec any, inputs []*graph.Node) (predictions []*graph.Node)\n",
    "```\n",
    "\n",
    "It takes a `context.Context` for the variables and hyperparameters, the `spec`  and a slice of `inputs` —\n",
    "the last two are fed by `Dataset.Yield` above. It returns a slice of `predictions` — is most cases there \n",
    "is just one value in the slice (only one prediction). During training `predictions` fed to the loss function,\n",
    "and during inference they can be returned directly.\n",
    "\n",
    "Our linear example has the simplest model possible:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "40eb46ee-4850-4c5f-8422-9273529a2c04",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import \"github.com/gomlx/gomlx/pkg/ml/context\"\n",
    "\n",
    "func modelGraph(ctx *context.Context, spec any, inputs []*Node) ([]*Node) {\n",
    "    _ = spec  // Not needed here, we know the dataset.\n",
    "    logits := layers.DenseWithBias(ctx, inputs[0], /* outputDim= */ 1)\n",
    "    return []*Node{logits}\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7c60a04-0eac-4ec9-8c91-888368802f30",
   "metadata": {},
   "source": [
    "> **Note**\n",
    "> - It uses the `layers.DenseWithBias` layer, which simply multiplies the input by a learnable matrix (weights)\n",
    ">   and add a learnable bias. It's the most basic building block of neural networks (NNs).\n",
    ">   The implementation of `layers.DenseWithBias` is pretty simple, and worth checking out to refresh how\n",
    ">   variables from the `Context` are used.\n",
    "> - Since it's a linear model, we don't use an activation function. The usual are available for NNs (`Relu`,\n",
    ">   `Sigmoid`, `Tanh` and more to come).\n",
    "> - The `spec` parameter allows the creation of a `ModelFn` that can be used for different types of data. The dataset\n",
    ">   can Yield also a `spec` about the type of data it is reading. Each different value of `spec` will trigger the\n",
    ">   the creation of a different computation graph, so ideally there would be at most a few types of different data\n",
    ">   source `spec`. Most commonly there is only one, like in this example, and the parameter can be ignored.\n",
    "\n",
    "### Trainer and Loop\n",
    "\n",
    "The last part is put together a `train.Trainer` and `train.Loop` objects in our `main()` function. The\n",
    "first stitches together the model, the optimizer and the loss function, and is able to run training\n",
    "steps and evaluations. The second, `train.Loop`, loops over the dataset executing a training step at\n",
    "a time and supporst a subscription (hooking) system, where one attaches things like a progress bar,\n",
    "or plotting of a graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7a7b1027-97dd-4ac6-8441-317c0999d598",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Target: coefficients=[0.0738 -0.915 -0.202], bias=-1.41\n",
      "Training data (inputs, labels): ((Float64)[10000 3], (Float64)[10000 1])\n",
      "\n",
      "      \u001b[1m 100% [========================================] (17333 steps/s)\u001b[0m [step=99] [loss+=0.0425] [~loss+=0.159] [~loss=0.159]                \n",
      "\n",
      "Learned: coefficients=[[0.0734] [-0.898] [-0.205]], bias=[-1.39]\n"
     ]
    }
   ],
   "source": [
    "import (\n",
    "    \"os\"\n",
    "    \"github.com/gomlx/gomlx/ui/commandline\"\n",
    "    \"github.com/gomlx/gomlx/pkg/ml/layers/regularizers\"\n",
    ")\n",
    "\n",
    "var (\n",
    "    flagNumExamples  = flag.Int(\"num_examples\", 10000, \"Number of examples to generate\")\n",
    "    flagNumFeatures  = flag.Int(\"num_features\", 3, \"Number of features\")\n",
    "    flagNoise        = flag.Float64(\"noise\", 0.2, \"Noise in synthetic data generation\")\n",
    "    flagNumSteps     = flag.Int(\"steps\", 100, \"Number of gradient descent steps to perform\")\n",
    "    flagLearningRate    = flag.Float64(\"lr\", 0.1, \"Initial learning rate.\")\n",
    ")\n",
    "\n",
    "// AttachToLoop decorators. It will be redefined later.\n",
    "func AttachToLoop(loop *train.Loop) {\n",
    "    commandline.AttachProgressBar(loop) // Attaches a progress bar to the loop.\n",
    "}\n",
    "\n",
    "// TrainMain() does everything to train the linear model.\n",
    "func TrainMain() {\n",
    "    flag.Parse()\n",
    "\n",
    "    // Select coefficients that we will try to predic.\n",
    "    trueCoefficients, trueBias := initCoefficients(backend, *flagNumFeatures)\n",
    "    fmt.Printf(\"Target: coefficients=%0.3v, bias=%0.3v\\n\", trueCoefficients.Value(), trueBias.Value())\n",
    "\n",
    "    // Generate training data with noise.\n",
    "    inputs, labels := buildExamples(backend, trueCoefficients, trueBias, *flagNumExamples, *flagNoise)\n",
    "    fmt.Printf(\"Training data (inputs, labels): (%s, %s)\\n\\n\", inputs.Shape(), labels.Shape())\n",
    "    dataset := &TrivialDataset{\"linear\", []*tensors.Tensor{inputs}, []*tensors.Tensor{labels}}\n",
    "\n",
    "    // Creates Context with learned weights and bias.\n",
    "    ctx := context.New()\n",
    "    ctx.SetParam(optimizers.ParamLearningRate, *flagLearningRate)  // = \"learning_rate\"\n",
    "    ctx.SetParam(regularizers.ParamL2, 1e-3)  // 1e-3 of L2 regularization.\n",
    "\n",
    "    // train.Trainer executes a training step.\n",
    "    trainer := train.NewTrainer(backend, ctx, modelGraph,\n",
    "        losses.MeanSquaredError,\n",
    "        optimizers.StochasticGradientDescent(),\n",
    "        nil, nil) // trainMetrics, evalMetrics\n",
    "    loop := train.NewLoop(trainer)\n",
    "    AttachToLoop(loop)\n",
    "\n",
    "    // Loop for given number of steps. must.M1() panics, if loop.RunSteps returns an error.\n",
    "    metrics := must(loop.RunSteps(dataset, *flagNumSteps))\n",
    "    _ = metrics  // We are not interested in them in this example.\n",
    "\n",
    "    // Print learned coefficients and bias -- from the weights in the dense layer.\n",
    "    fmt.Println()\n",
    "    coefVar, biasVar := ctx.GetVariableByScopeAndName(\"/dense\", \"weights\"), ctx.GetVariableByScopeAndName(\"/dense\", \"biases\")\n",
    "    learnedCoef, learnedBias := must(coefVar.Value()), must(biasVar.Value())\n",
    "    fmt.Printf(\"Learned: coefficients=%0.3v, bias=%0.3v\\n\", learnedCoef.Value(), learnedBias.Value())\n",
    "}\n",
    "\n",
    "%%\n",
    "TrainMain()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ac53dff-00d2-4372-815a-fa8ee78f2572",
   "metadata": {},
   "source": [
    "> **Note**:\n",
    "> - Hyperparameters are set on the context. Layers and optimizers can define their own\n",
    ">   hyperparemeters independently. `Context` uses a scoping mechanism (like directories), \n",
    ">   and hyperparameters can take specialized values under specific scopes -- e.g.: \n",
    ">   doing `ctx.In(\"dense_5\").SetParam(regularizers.ParamL2, 0.1)` would set\n",
    ">   L2 regularization only for the layer `dense_5` of the model to `0.1`.\n",
    "> - The `trainer` constructor also takes as input arbitrary metrics (for training and evaluation).\n",
    ">   The metric of the loss of the last batch, and a moving average of the loss are always included\n",
    ">   automatically. There are many others, that can be means or moving averages, etc.\n",
    "> - The 'Loop' object is very flexible. One can attach any functionality (hooks) to be executed\n",
    ">   during training with `OnStart`, `OnStep`, `OnEnd`, `EveryNSteps`, `NTimesDuringLoop`, \n",
    ">   `PeriodicCallback` (e.g.: every N minutes) and `ExponentialCallback`.\n",
    "> - The most common such functionality is the `commandline.AttachProgressBar`. There is also a plotting\n",
    ">   of any arbitrary metric or any arbitrary `Node` in the computation graph.\n",
    "> - `Loop.RunSteps()` returns also the final metrics from the training, usually printed out.\n",
    "\n",
    "### Training and Plotting\n",
    "\n",
    "As our last example, let's train it \"for real\", that is, with more steps.\n",
    "\n",
    "And to make things prettier, let's attach also a plot of the metrics registered. In our example\n",
    "the only metrics are the default ones, the batch and mean of the loss -- the mean squared error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e6514ede-c84e-4259-96ee-8d6eea196815",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Target: coefficients=[-3.21 0.609 -6.14], bias=3.84\n",
      "Training data (inputs, labels): ((Float64)[10000 3], (Float64)[10000 1])\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m   3% [>.......................................] (12275 steps/s) [0s:0s]\u001b[0m [step=164] [loss+=0.0926] [~loss+=2.06] [~loss=2.02]        "
     ]
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m 100% [========================================] (15646 steps/s)\u001b[0m [step=4999] [loss+=0.0883] [~loss+=0.0883] [~loss=0.04]                  \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<p><b>Metric: loss</b></p>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div id=\"7b67ca46\"></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script charset=\"UTF-8\">\n",
       "(() => {\n",
       "\tconst src=\"https://cdn.plot.ly/plotly-2.34.0.min.js\";\n",
       "\tvar runJSFn = function(module) {\n",
       "\t\t\n",
       "\tif (!module) {\n",
       "\t\tmodule = window.Plotly;\n",
       "\t}\n",
       "\tlet data = JSON.parse('{\"data\":[{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Batch Loss+Regularization\",\"x\":[51,106,167,234,308,389,478,576,684,803,934,1078,1236,1410,1601,1811,2042,2296,2575,2882,3220,3592,4001,4451,4946,5000],\"y\":[0.49201912743406506,0.12099068650598462,0.09233659096595803,0.08887947850278159,0.0883599266590213,0.08827325475447512,0.08825770328219987,0.08825481856845302,0.08825427922052473,0.08825417929935148,0.08825416117279873,0.08825415797982257,0.08825415743694982,0.08825415734807174,0.08825415733428105,0.0882541573322501,0.08825415733196892,0.08825415733193256,0.08825415733192822,0.08825415733192772,0.08825415733192768,0.08825415733192767,0.08825415733192767,0.08825415733192767,0.08825415733192767,0.08825415733192767]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Loss+Regularization\",\"x\":[51,106,167,234,308,389,478,576,684,803,934,1078,1236,1410,1601,1811,2042,2296,2575,2882,3220,3592,4001,4451,4946,5000],\"y\":[7.348469653472045,3.6437257341507103,2.0195848352642303,1.0739884258133408,0.5569465418216063,0.29593187321454023,0.17316171236682623,0.1199650163969135,0.09896481776933402,0.09149314233188514,0.08912236282547768,0.08845837479532227,0.08829588801999258,0.08826141807677187,0.08825522222294381,0.08825428636414241,0.08825416999184715,0.0882541583176958,0.08825415739163128,0.08825415733465708,0.08825415733201944,0.08825415733193023,0.08825415733192876,0.08825415733192876,0.08825415733192876,0.08825415733192876]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Loss\",\"x\":[51,106,167,234,308,389,478,576,684,803,934,1078,1236,1410,1601,1811,2042,2296,2575,2882,3220,3592,4001,4451,4946,5000],\"y\":[7.320463913503426,3.607387263348591,1.9783791947914087,1.029563847943385,0.5106107003247292,0.24856782190253618,0.12528624961898016,0.07185542057102069,0.050757079172817354,0.04324788548741037,0.04086403116402077,0.04019586331054854,0.0400321291889823,0.03999729920000741,0.039990999253930834,0.03999003190322041,0.03998990550826277,0.03998989055309237,0.03998988855352181,0.03998988815181851,0.03998988804209222,0.03998988801007331,0.03998988800097511,0.039989887998503594,0.039989887997866604,0.03998988799783844]}],\"layout\":{\"legend\":{},\"title\":{\"text\":\"loss\"},\"xaxis\":{\"showgrid\":true,\"title\":{\"text\":\"Steps\"},\"type\":\"log\"},\"yaxis\":{\"showgrid\":true,\"type\":\"log\"}}}');\n",
       "\tmodule.newPlot('7b67ca46', data);\n",
       "\n",
       "\t}\n",
       "\t\n",
       "    if (typeof requirejs === \"function\") {\n",
       "        // Use RequireJS to load module.\n",
       "\t\tlet srcWithoutExtension = src.substring(0, src.lastIndexOf(\".js\"));\n",
       "        requirejs.config({\n",
       "            paths: {\n",
       "                'plotly': srcWithoutExtension\n",
       "            }\n",
       "        });\n",
       "        require(['plotly'], function(plotly) {\n",
       "            runJSFn(plotly)\n",
       "        });\n",
       "        return\n",
       "    }\n",
       "\n",
       "\tvar currentScripts = document.head.getElementsByTagName(\"script\");\n",
       "\tfor (const idx in currentScripts) {\n",
       "\t\tlet script = currentScripts[idx];\n",
       "\t\tif (script.src == src) {\n",
       "\t\t\trunJSFn(null);\n",
       "\t\t\treturn;\n",
       "\t\t}\n",
       "\t}\n",
       "\n",
       "\tvar script = document.createElement(\"script\");\n",
       "\n",
       "\tscript.charset = \"utf-8\";\n",
       "\t\n",
       "\tscript.src = src;\n",
       "\tscript.onload = script.onreadystatechange = function () { runJSFn(null); };\n",
       "\tdocument.head.appendChild(script);\t\n",
       "})();\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Learned: coefficients=[[-3.2] [0.611] [-6.13]], bias=[3.84]\n"
     ]
    }
   ],
   "source": [
    "import \"github.com/gomlx/gomlx/ui/gonb/plotly\"\n",
    "import \"os\"\n",
    "\n",
    "func AttachToLoop(loop *train.Loop) {\n",
    "    commandline.AttachProgressBar(loop) // Attaches a progress bar to the loop.\n",
    "    _ = plotly.New().Dynamic().ScheduleExponential(loop, 50, 1.1)\n",
    "}\n",
    "\n",
    "%% --steps=5000\n",
    "TrainMain()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8268a8d3-148d-4db5-9289-573d36b73412",
   "metadata": {},
   "source": [
    "> **Note**:\n",
    "> - The `gomlx/examples/notebook/gonb/plotly` package will automatically plot all the metrics registered in the trainer. When it attaches itself to \n",
    "> `loop` it collects the metric values (and run evaluation on any requested datasets), and at the end of it, plot it.\n",
    "> - Optionally, with `Plots.Dynamic()` it also plots intermediary results, displaying the progress of the training as it happens.\n",
    "> - Previously, we added a little L2 regularization as a hyperparameter. The `layers.DenseWithBias` layer picks up on that hyperparameter, and adds that bit of L2 regularization. GoMLX tracks the loss with/without regularization separately, and one can see the difference in the two lines.\n",
    "\n",
    "### Other examples\n",
    "\n",
    "There is much more in the libraries and [examples](https://github.com/gomlx/gomlx/tree/main/examples). \n",
    "GoMLX libraries are well documented, and the implementation is generally simple to use and understand -- don't hesitate into diving into the code. \n",
    "For machine learning, we highly recommend diving into the `context.Context` package/objects, and looking at the `layers` library, to learn how many of the common ML layers work.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a71408a0-1763-4ebe-b7f7-912de9866863",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Debugging\n",
    "\n",
    "Unfortunately, the computers just \"don't get it\": they do exactly what we tell them to do,\n",
    "as opposed to what we want them to do, and thus programs fail or crash.\n",
    "And AIs sometimes \"get it\", but still write wrong programs.\n",
    "\n",
    "GoMLX provides different ways to track down various types of errors. The most commonly used below:\n",
    "\n",
    "### Good old \"printf\"\n",
    "\n",
    "It's convenient because of Go fast compilation (change something and run to see what one gets is\n",
    "almost instant). Logging results to stdout is a valid way of developing. During graph building development,\n",
    "often one prints the shape of the `Node` being operated to confirm (or not) one's expectations.\n",
    "\n",
    "### Errors in **GoMLX** have stacktrace\n",
    "\n",
    "Errors during the building of the graph are reported back with `panic`. \n",
    "The `graph.Exec` or `context.Exec` recover and convert them to errors for you.\n",
    "\n",
    "These errors include a stack-trace the tells you exactly where things happened. The stack-grace can be printed with `\"%+v\"` format string.\n",
    "\n",
    "### Node Shape Asserts\n",
    "\n",
    "During the writing of complex models, it's very common to add comments on the expected shapes of the graph nodes, to\n",
    "facilitate the reader (and developer) of the code to have the right mental model of what is going on.\n",
    "\n",
    "GoMLX provides a series of _assert_ methods that can be used instead. They serve both as documentation, and an early\n",
    "exit in case of some unexpected results. \n",
    "\n",
    "For example, a `modelGraph` function could contain:\n",
    "\n",
    "```go\n",
    "    batch_size := inputs[0].Shape().Dimensions[0]\n",
    "    ...\n",
    "    layer := Concatenate(allEmbeddings, -1)\n",
    "    layer.AssertDims(batchSize, -1)  // 2D tensor, with batch size as the leading dimension.\n",
    "```\n",
    "\n",
    "Although using these when building graphs is the most common case, there are similar assert functions for tensors and shapes themselves in the package `gomlx/types/shapes`.\n",
    "\n",
    "\n",
    "### Graph Execution Logging\n",
    "\n",
    "Every `Node` of the graph can be flagged with `SetLogged(msg)`.\n",
    "The executor (`Exec`) will at the end of the execution log all these values. The default logger \n",
    "(set with `Exec.SetLogger`) will simply print the message `msg` along with the value of the `Node` of\n",
    "interest. Creating a specialized loggers that handle arbitrary nodes is trivial.\n",
    "\n",
    "### Catching `NaN` and `Inf` in your training\n",
    "\n",
    "This is a common source of headaches when training complex models. \n",
    "The package `nanlogger` implements a specialized logger and monitor results on arbitrary nodes in the graph,\n",
    "and will `panic` with custom messages (and a stack-trace) a the first sight of a `NaN` (or `±Inf`).\n",
    "\n",
    "### More Debugging\n",
    "\n",
    "Tests and these methods have been enough to develop most of GoMLX so far. \n",
    "But there are other debugging tools that could be made, see discussion in the [Debugging](https://github.com/gomlx/gomlx/blob/main/docs/debugging.md) document. Let us know if you need something specialized.\n",
    "\n",
    "---\n",
    "\n",
    "Happy coding and [good luck](https://arxiv.org/abs/1803.03635) on modeling!!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeb81912-1bcb-4e28-a397-52e96b20744b",
   "metadata": {
    "tags": []
   },
   "source": [
    "<img alt=\"[Zürich See]\" src=\"zurich_see.jpg\" style=\"width:100%;\"/>\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Go (gonb)",
   "language": "go",
   "name": "gonb"
  },
  "language_info": {
   "codemirror_mode": "",
   "file_extension": ".go",
   "mimetype": "text/x-go",
   "name": "go",
   "nbconvert_exporter": "",
   "pygments_lexer": "",
   "version": "go1.25.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
